//
//  main.swift
//  DataGenerator
//
//  Created by AMalashin on 25.09.2023.
//
// The utility to generate a custom LM trining data
// ~/Library/Developer/Xcode/DerivedData/DentalFormula-fcbytnozlnuhfdgxcbnipipixgmc/Build/Products/Debug/DataGenerator

import Foundation
import Speech

enum SupportedLanguages {
    case russian
    case english
    
    var id: String {
        switch self {
        case .russian: return "ru_RU"
        case .english: return "en_US"
        }
    }
}

struct DefectLocalized {
    let language: SupportedLanguages
    let translations: [String]
    let phrase: String
}

let defectsInRU = DefectLocalized(language: .russian, translations: ["кариес", "пломба", "удалён", "кариес-пломба", "пломба-кариес", "интактный"], phrase: "Зуб <number> <defect>")
let defectsInEN = DefectLocalized(language: .english, translations: ["decay", "filling", "missing", "decay-filling", "filling-decay", "intact"], phrase: "Tooth <number> <defect>")

func generateCustomLMData(defect: DefectLocalized) async /*-> SFCustomLanguageModelData*/ {
    var numbers: [String] = []
    for q in 1...4 {
        for n in 1...8 {
            numbers.append("\(q)\(n)")
        }
    }

    let data = SFCustomLanguageModelData(locale: Locale(identifier: defect.language.id), identifier: "in.malash.DentalFormula", version: "1.0") {
//        SFCustomLanguageModelData.PhraseCount(phrase: "Зуб 17 кариес пломба", count: 10)
        
        if defect.language == .russian {
            SFCustomLanguageModelData.CustomPronunciation(grapheme: "интактный", phonemes: ["I n t \"a k t n \"1 j"])
            SFCustomLanguageModelData.PhraseCount(phrase: "Зуб интактный", count: 10)
        }
        
        SFCustomLanguageModelData.PhraseCountsFromTemplates(classes: [
            "number": numbers,
            "defect": defect.translations
        ]) {
            SFCustomLanguageModelData.TemplatePhraseCountGenerator.Template(defect.phrase, count: 10_000)
        }

    }

    do {
        let fileName = "/var/tmp/lm-data-\(defect.language)"
        try await data.export(to: URL(filePath: fileName))
        print("LM Data for \(defect.language) saved to \(fileName)")
    } catch {
        print("Error exporting LM data \(error.localizedDescription)")
    }
//    return data
}

await generateCustomLMData(defect: defectsInEN)
await generateCustomLMData(defect: defectsInRU)
